
import networkx as nx
import numpy as np
from pyteomics import pylab_aux,mass
import pylab

from lib import *


def add_parent_by(data):
    temp=read_json("amino_acid.json")
    Water_weight=temp[1]["Monoisotopic Mass"]
    H_weight=temp[2]["Monoisotopic Mass"]/2
    mz=data["precursor_mz"]
    intensity=data["precursor_intensity"]/2
    charge=data["precursor_charge"]
    data["product_intensity"].extend([intensity]*2)
    data["product_mz"].append(mz*charge-(charge-1)*H_weight-Water_weight)
    data["product_mz"].append(mz*charge-(charge-1)*H_weight)
    return data

def __get_spectrum(DATA,title:str=None):

    spectrum={
        'params':{
            'title':title,
            'pepmass':(390.70524061033,None),
            'charge':[2]
        },
        'm/z array':np.array(DATA['product_mz']),
        'intensity array':np.array(DATA['product_intensity'])
    }
    return spectrum

def plot_spectrum(DATA):
    spectrum=__get_spectrum(DATA)
    pylab.figure()
    pylab_aux.plot_spectrum(spectrum, title="Experimental spectrum ")
    pylab.show()

def plot_spectrum_icon(DATA,peptide):
    spectrum=__get_spectrum(DATA)
    pylab.figure()
    pylab_aux.annotate_spectrum(spectrum, peptide,
    title='Annotated spectrum '+ peptide,
    maxcharge=spectrum['params']['charge'][0],ion_types=('a','b','c','x','y','z'))
    pylab.show()

def cleaning(DATA,min_intensity):
    DATA['product_mz']=np.array(DATA['product_mz'])
    DATA['product_intensity']=np.array(DATA['product_intensity'])
    clear=DATA['product_intensity']>min_intensity
    DATA['product_mz']=DATA['product_mz'][clear]
    p_i=DATA['product_intensity'][clear]
    pii=list(enumerate(p_i))
    pii.sort(key= lambda x : x[1])
    p_i=list(enumerate(pii))
    p_i.sort(key= lambda x : x[1][0])
    DATA['product_intensity_normal']=np.array([(pi[0]+1)/len(p_i) for pi in p_i])
    DATA['product_intensity']=[pi[1][1] for pi in p_i]
    return DATA

def decoding(DATA):

    amino_acids_Letter=[ aa["Letter"] for aa in get_all_ppfms() ]
    amino_acids_Mass=[ aa["Mass"] for aa in get_all_ppfms() ]
    G = nx.DiGraph()

    N_samples=len(DATA["product_mz"])
    for i in range(N_samples):
        for j in range(N_samples):
            fpp=DATA["product_mz"][i]
            tpp=DATA["product_mz"][j]
            if tpp <= fpp :
                continue
            dp = tpp - fpp
            for index_aa,aa in enumerate(amino_acids_Mass):
                if abs(dp-aa)/aa<=20E-6:
                    G.add_edges_from(
                        [(i,j),],
                        weight=(
                            DATA["product_intensity"][j]+
                            DATA["product_intensity"][j]
                            )/2,
                        weight_normal=(
                            DATA["product_intensity_normal"][j]+
                            DATA["product_intensity_normal"][j]
                            )/2,
                        aa_id=index_aa,
                        aa_mass=aa
                        )

    over = False
    while not over:
        longest_path = nx.dag_longest_path(G,weight="weight")
        index_data=longest_path[0]
        first_mz=DATA["product_mz"][index_data]
        for icon in get_all_icon():
            true_mz=icon["m/z"]
            if abs(first_mz-true_mz)/true_mz>20E-6:
                continue

            print(f"- {icon['ppf']}")
            print(f"m/z: {first_mz:0.6f}, deltaP_mz: {abs(true_mz-first_mz):0.6f}")
            print(f"type:{icon['type']}")
            print(f"product_mz: {DATA['product_mz'][index_data]}")
            print(f"product_intensity: {DATA['product_intensity'][index_data]:0.6f}, is {DATA['product_intensity_normal'][index_data]:0.6f}")
            over=True
            peptide_seq=icon['ppf']
            is_b_icon = True if icon['type'][0]=="b" else False
            break

        else:
            G.remove_edge(index_data,longest_path[1])
    
    mass=first_mz
    for index_aa,index_data in enumerate(longest_path):
        if index_aa == 0:
            continue
        print()
        mass=DATA["product_mz"][index_data]-mass
        aa_id=G.get_edge_data(longest_path[index_aa-1],index_data)["aa_id"]
        print("+",amino_acids_Letter[aa_id])
        print(f"mass: {mass:0.6f}, deltaP_mass: {abs(mass-amino_acids_Mass[aa_id]):0.6f}")
        print(f"product_mz: {DATA['product_mz'][index_data]}")
        print(f"product_intensity: {DATA['product_intensity'][index_data]:0.6f}, is {DATA['product_intensity_normal'][index_data]:0.6f}")
        mass=DATA["product_mz"][index_data]
        peptide_seq+=amino_acids_Letter[aa_id]
    
    return peptide_seq if is_b_icon else peptide_seq[::-1]


def main(min_intensity:float=5000):
    DATA=read_json("data.json")

    plot_spectrum(DATA)

    DATA=add_parent_by(DATA)
    DATA=cleaning(DATA,min_intensity)
    peptide_seq=decoding(DATA)

    plot_spectrum_icon(read_json("data.json"),peptide_seq)


if __name__ == "__main__":
    main()